#!/usr/bin/env python
#
# Author: LHC simple program
#

'''
A simple example to run RHF 

'''

from pyscf import scf,gto,dft
import pyscf
import numpy as np
import scipy
from functools import partial
import warnings

#ol = pyscf.M(
#   atom = 'H 0 0 0; F 0 0 1.1',  # in Angstrom
#   basis = 'sto-3g',
#   symmetry = True,
#
mol = gto.Mole()
mol.atom = '''
O        0.000000    0.000000    0.117790
H        0.000000    0.755453   -0.471161
H        0.000000   -0.755453   -0.471161'''
mol.basis = "6-31G"
mol.verbose = 3
mol.build()
  ### unit :bohr
#pyscf interface
A_t=mol.atom_coords() ###输出unit:A
#print(A_t)
T = mol.intor("int1e_kin")
V = mol.intor('int1e_nuc')
S = mol.intor('int1e_ovlp')
X = scipy.linalg.fractional_matrix_power(S,-0.5)##对S矩阵求S^(-1/2)
eri=mol.intor('int2e')
#print(eri.shape)
Hcore=T + V

scf_eng=dft.RKS(mol)
scf_eng.xc='Slater'
#定义格点
grids=dft.Grids(mol) #生成格点
grids.atom_grid=(99,590)
#grids.build()
#格点数量
#Becke 网格生成方案,更好计算电子密度
grids.becke_scheme = dft.gen_grid.stratmann
grids.build()
ngrid=grids.weights.size
#print(ngrid)
ni = dft.numint.NumInt()

weight=grids.weights#仅用于与泛函格点乘积，在公式中不会出现
#print(weight.shape)

#print(grids)
natm=mol.natm    #atom number
nmo =nao=mol.nao #atom orbital number
nocc=mol.nelec[0] #[0]:Alpha ele [1]:Beta ele
#print(nocc)
so,sv,sa=slice(0,nocc),slice(nocc,nmo),slice(0,nmo)  # so:(0,occ-1);sv:(occ,nmo-1)
#print(so)

#####获得笛卡尔坐标，核电荷数，E_nuc###
A_t = mol.atom_coords()
#print(A_t)
r_ABt =A_t[:,None,:] - A_t[None,:,:]
#print(A_t[:,None,:])
#print("=======")
#print(A_t[None,:,:])
r_AB = np.linalg.norm(r_ABt,axis=-1)
r_AB += np.diag(np.ones(mol.natm)*np.inf)
A_charge=mol.atom_charges()[:,None]*mol.atom_charges()
#print(A_charge)
E_nuc=0.5 *(A_charge / r_AB).sum()
print('Neucleus energy   ',E_nuc,'a.u.')
###Provide a first guess using pyscf
mf=scf.RHF(mol)
#D=mf.get_init_guess(mol,key='minao') # 1)
#D=mf.get_init_guess(mol,key='huckel')# 2)
#D=mf.get_init_guess(mol,key='atom')  # 3)
#D=mf.get_init_guess(mol,key='1e')    # 4)
###end guess###

ao = np.zeros((4, ngrid, nao))  # 4 refers to (noderiv, x_deriv, y_deriv, z_deriv)
				#这里需要注意ao的维度，如果等于1，则为4，等于2，则为10
g_start = 0
for inner_ao, _, _, _ in ni.block_loop(mol, grids, nao, deriv=1, max_memory=2000):
    ao[:, g_start:g_start+inner_ao.shape[-2]] = inner_ao
    g_start += inner_ao.shape[-2]
ao_0 = ao[0]
ao_1 = ao[1:4]


def gen_J(R):
    return np.einsum("uvkl, kl -> uv", eri, R)

def gen_K(R):
    return np.einsum("ukvl, kl -> uv", eri, R)

def gen_rho_grid(D):#D->Density
    rho_0 = np.einsum("gu, gv, uv -> g", ao_0, ao_0, D)      #密度
    rho_1 = 2 * np.einsum("rgu, gv, uv -> rg", ao_1, ao_0, D)#密度梯度
    rho_01 = np.concatenate(([rho_0], rho_1), axis=0)        #密度和密度梯度的合并张量
    #print(rho_0.shape,rho_01.shape)
    return rho_0, rho_1, rho_01

def gen_kernel_grid(xc_code, rho_0):
    cx = ni.hybrid_coeff(xc_code)  #
    tmp = ni.eval_xc(xc_code, rho_0, deriv=1)
    exc, vxc = tmp[:2]
    fr=vxc[0]
    #exc = np.asarray(exc).reshape(ngrid,)
    #vxc = np.asarray(vxc).reshape(ngrid,)
    #print('exc shape:',exc.shape)
    exc *= weight
    fr *= weight
    #fg *= weight
    return cx, exc, fr

def gen_e_slater(D):
    rho_0, rho_1, rho_01 = gen_rho_grid(D)
    kernel_s=-3/4 * (3/np.pi)**(1/3)*rho_0**(1/3)
    Exc = (grids.weights * kernel_s * rho_0).sum()
    return Exc

def gen_vxc(xc_code, D):#交换相关势
    rho_0, rho_1, rho_01 = gen_rho_grid(D)
    _, _, fr = gen_kernel_grid(xc_code, rho_0)
    vxc = (
         np.einsum("g, gu, gv -> uv", fr, ao_0, ao_0)
    )
    return vxc
 

def gen_F_0_ao(xc_code, D):
    cx = ni.hybrid_coeff(xc_code)
    F_0_ao = Hcore + gen_J(D) - 0.5 * cx * gen_K(D)+ gen_vxc(xc_code, D) #如果是纯泛函则(cx=0) 
    return F_0_ao

def gen_energy_elec(xc_code, D):
    cx = ni.hybrid_coeff(xc_code)
    rho_0, _, rho_01 = gen_rho_grid(D)           #return rho_0,rho_1,rho_01(密度，密度梯度，密度和密度梯度合并张量)
    cx, exc, fr = gen_kernel_grid(xc_code, rho_0)#return cx,exc,fr,fg
    energy_elec = ((Hcore + 0.5 * gen_J(D)) * D).sum()
    energy_elec += gen_e_slater(D)
    return energy_elec


#C = e = NotImplemented
D = np.zeros((nao, nao))
D_old = np.zeros((nao, nao)) + 1e-4
count = 0

while (not np.allclose(D, D_old, atol=1e-8, rtol=1e-6)):
    if count > 500:
        raise ValueError("SCF not converged!")
    count += 1
    D_old = D
    F_0_ao = gen_F_0_ao("slater", D)
    e, C = scipy.linalg.eigh(F_0_ao, S)  # Solve FC = SCe
    D = 2 * C[:, so] @ C[:, so].T
    if count > 1:
        D = 0.3 * D + 0.7 * D_old             # For convergence(damping)

E_elec = gen_energy_elec("slater", D)
E_tot = E_elec + E_nuc
E_xc=gen_e_slater(D)
print(E_xc)

#rho_0,_,_= gen_rho_grid(D)
#print(rho_0.shape)
#print((rho_0*grids.weights).sum())
print("SCF Converged in          ", count, " loops")
print("Electronic energy         ", E_elec, " a.u.")
print("Total energy              ", E_tot, " a.u.")





